package _back_tracking;

import org.junit.Assert;
import org.junit.Test;

import java.util.Arrays;

/**
 * TODO: 698. 划分为k个相等的子集
 * https://blog.51cto.com/acoier/5708749
 * https://leetcode.cn/problems/partition-to-k-equal-sum-subsets/description/
 * <p>
 * ```
 * 698. 划分为k个相等的子集
 * 给定一个整数数组  nums 和一个正整数 k，找出是否有可能把这个数组分成 k 个非空子集，其总和都相等。
 * <p>
 * 示例 1：
 * 输入： nums = [4, 3, 2, 3, 5, 2, 1], k = 4
 * 输出： True
 * 说明： 有可能将其分成 4 个子集（5），（1,4），（2,3），（2,3）等于总和。
 * <p>
 * 示例 2:
 * 输入: nums = [1,2,3,4], k = 3
 * 输出: false
 * <p>
 * <p>
 * 提示：
 * 1 <= k <= len(nums) <= 16
 * 0 < nums[i] < 10000
 * 每个元素的频率在 [1,4] 范围内
 * ```
 */
public class leetcode_698 {
    @Test
    public void test_1() {
        boolean result = canPartitionKSubsets(new int[]{1, 2, 3, 4}, k = 4);
        boolean expected = true;
        Assert.assertEquals(expected, result);
    }


    @Test
    public void test_2() {
        boolean result = canPartitionKSubsets(new int[]{1, 2, 3, 4}, k = 3);
        boolean expected = false;
        Assert.assertEquals(expected, result);
    }

    int[] nums;
    int size, targetSum, k;

    public boolean canPartitionKSubsets(int[] nums, int k) {
        this.nums = nums;
        this.k = k;
        int sum = 0;
        for (int x : this.nums) sum += x;
        if (sum % this.k != 0) return false; // 可行性剪枝
        Arrays.sort(this.nums);
        size = this.nums.length;
        targetSum = sum / this.k;
        return dfs(size - 1, 0, 0, new boolean[size]);
    }

    boolean dfs(int index, int currentSum, int currentK, boolean[] used) {
        if (currentK == k) return true;
        if (currentSum == targetSum) {
            return dfs(size - 1, 0, currentK + 1, used);
        }
        for (int i = index; i >= 0; i--) {  // 顺序性剪枝
            if (used[i] || currentSum + nums[i] > targetSum) continue;  // 可行性剪枝
            used[i] = true;
            if (dfs(i - 1, currentSum + nums[i], currentK, used)) return true;
            used[i] = false;
            if (currentSum == 0) return false; // 可行性剪枝
        }
        return false;
    }
}